package weeny.pre.io;

import java.awt.Font;
import java.awt.geom.Rectangle2D;
import java.io.OutputStream;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;

import prefuse.data.Edge;
import prefuse.data.Graph;
import prefuse.data.Node;
import prefuse.data.Table;
import prefuse.data.column.Column;
import prefuse.data.io.AbstractGraphWriter;
import prefuse.data.io.DataIOException;
import prefuse.data.parser.CollectionParser;
import prefuse.data.parser.DataParser;
import prefuse.data.parser.ParserFactory;
import prefuse.util.FontLib;
import prefuse.util.io.XMLWriter;
import weeny.pre.display.Initialization;
import weeny.util.io.IOUtil;

/**
 * This writer supports several types of columns, for more details, please see
 * {@link MyGraphReader}
 * 
 * @author weeny
 *
 */
public class MyGraphWriter extends AbstractGraphWriter{
	
    /**
     * String tokens used in the GraphML format.
     */
    public static interface Tokens {
        public static final String ID            = "id";
        public static final String GRAPH         = "graph";
        public static final String EDGEDEF       = "edgedefault";
        public static final String DIRECTED      = "directed";
        public static final String UNDIRECTED    = "undirected";
        public static final String NODEKEY       = "nodekey";
        
        
        public static final String KEY           = "key";
        public static final String FOR           = "for";
        public static final String ALL           = "all";
        public static final String ATTRNAME      = "attr.name";
        public static final String ATTRTYPE      = "attr.type";
        public static final String DEFAULT       = "default";
        
        public static final String NODE          = "node";
        public static final String EDGE          = "edge";
        public static final String SOURCE        = "source";
        public static final String TARGET        = "target";
        public static final String DATA          = "data";
        public static final String TYPE          = "type";
        
        public static final String INT           = "int";
        public static final String INTEGER       = "integer";
        public static final String LONG          = "long";
        public static final String FLOAT         = "float";
        public static final String DOUBLE        = "double";
        public static final String REAL          = "real";
        public static final String BOOLEAN       = "boolean";
        public static final String STRING        = "string";
        public static final String INT_ARRAY     = "intarray";
        public static final String FLOAT_ARRAY   = "floatarray";
        public static final String DOUBLE_ARRAY  = "doublearray";
        public static final String BOOLEAN_ARRAY = "booleanarray";
        public static final String LONG_ARRAY    = "longarray";
        public static final String STRING_ARRAY  = "stringarray";

        public static final String GRAPHML = "graphml";
        public static final String GRAPHML_HEADER =
            "<graphml xmlns=\"http://graphml.graphdrawing.org/xmlns\"\n" 
            +"  xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n"
            +"  xsi:schemaLocation=\"http://graphml.graphdrawing.org/xmlns\n"
            +"  http://graphml.graphdrawing.org/xmlns/1.0/graphml.xsd\">\n\n";
    }
    
    //just for convenient
    private static final HashMap<Class, String> TYPES = new HashMap();
    static {
        TYPES.put(int.class,       Tokens.INT);
        TYPES.put(long.class,      Tokens.LONG);
        TYPES.put(float.class,     Tokens.FLOAT);
        TYPES.put(double.class,    Tokens.DOUBLE);
        TYPES.put(boolean.class,   Tokens.BOOLEAN);
        TYPES.put(String.class,    Tokens.STRING);
        TYPES.put(int[].class,     Tokens.INT_ARRAY);
        TYPES.put(float[].class,   Tokens.FLOAT_ARRAY);
        TYPES.put(double[].class,  Tokens.DOUBLE_ARRAY);
        TYPES.put(boolean[].class, Tokens.BOOLEAN_ARRAY);
        TYPES.put(long[].class,    Tokens.LONG_ARRAY);
        TYPES.put(String[].class,  Tokens.STRING_ARRAY);
    }
    
	public static void main(String[] args) throws Exception {
		Initialization.init();
		Graph g = new Graph();
		g.getNodeTable().addColumn("A", HashSet.class, null);
		g.getNodeTable().addRow();
		HashSet<Integer> set = new HashSet<Integer>();
		set.add(11);
		set.add(23);
		g.getNode(0).set("A", set);
		g.getNodeTable().addColumn("B", Font.class, null);
		g.getNode(0).set("B", FontLib.getFont("Arial", 10));
		g.getNodeTable().addColumn("C", Rectangle2D.class, null);
		g.getNode(0).set("C", new Rectangle2D.Double());
		g.getNodeTable().addColumn("D", ArrayList.class, null);
		ArrayList<String> list = new ArrayList<String>();
		list.add("Q");
		list.add("W,W");
		list.add("E\"E");
		g.getNode(0).set("D", list);
		g.getNodeTable().addColumn("E", String[].class, null);
		g.getNode(0).set("E", new String[]{"X","Y,Y","Z\"Z"});
		new MyGraphWriter().writeGraph(g, IOUtil.desktop()+"A.xml");
	}
	
	//Remember all the old parsers, because we may need to change the parsers
	//during the writing process. So we need to restore them back after the
	//graph is saved~
	protected Map<String, DataParser> cacheParsers(Table table){
		Map<String, DataParser> oldParsers = new HashMap();
		for(int i = 0; i<table.getColumnCount(); i++){
			String field = table.getColumnName(i);
			oldParsers.put(field, table.getColumn(i).getParser());
		}
		return oldParsers;
	}
	//put all the old parsers back to the table
	protected void restoreParsers(Table table, Map<String, DataParser> map){
		for(int i = 0; i<table.getColumnCount(); i++){
			String field = table.getColumnName(i);
			table.getColumn(i).setParser(map.get(field));
		}
	}
	
	protected void checkAndPrepareParsers(Table table){
		for(int col = 0; col<table.getColumnCount(); col++){
			Class type = table.getColumnType(col);
			DataParser dp = table.getColumn(col).getParser();
			//if the CollectionParser is not ready
			//try to format elements until it is ready
			if (dp instanceof CollectionParser) {
				for (int row = 0; row < table.getRowCount(); row++){
					if (((CollectionParser)dp).isReady()) {
						break;
					}
					((CollectionParser)dp).init(table.get(row, col));
				}
			}
			else {
				dp = ParserFactory.getDefaultFactory().getParser(type);
			}
			
			if (dp == null){
				throw new IllegalStateException("Cannot find parser for "+type);
			}
			table.getColumn(col).setParser(dp);
		}
	}
	public void safeWrite(Graph graph, String file){
		try {
			writeGraph(graph, file);
		} catch (DataIOException e) {
			e.printStackTrace();
		}
	}
	public void safeWrite(Graph graph, OutputStream os){
		try {
			writeGraph(graph, os);
		} catch (DataIOException e) {
			e.printStackTrace();
		}
	}
	public void writeGraph(Graph g, OutputStream os) throws DataIOException {
		Map<String, DataParser> nodeOldParser = cacheParsers(g.getNodeTable());
		Map<String, DataParser> edgeOldParser = cacheParsers(g.getEdgeTable());
		checkAndPrepareParsers(g.getNodeTable());
		checkAndPrepareParsers(g.getEdgeTable());
		
        XMLWriter xml = new XMLWriter(new PrintWriter(os));
        xml.begin(Tokens.GRAPHML_HEADER, 2);
        
        xml.comment("GraphML Writer: " + new Date(System.currentTimeMillis()));

        // print the graph schema
        printSchema(xml, Tokens.NODE, g.getNodeTable(), null);
        printSchema(xml, Tokens.EDGE, g.getEdgeTable(), 
        		g.getNodeKeyField()!= null? null :
        		new String[]{g.getEdgeSourceField(), g.getEdgeTargetField()});
        xml.println();
        
        // print graph contents
        if (g.getNodeKeyField() == null) {
        	xml.start(Tokens.GRAPH, Tokens.EDGEDEF,
        		g.isDirected() ? Tokens.DIRECTED : Tokens.UNDIRECTED);
        }
        else {
        	String[] attr = new String[]{Tokens.EDGEDEF, Tokens.NODEKEY};
        	String[] vals = new String[attr.length];
        	vals[0] = g.isDirected() ? Tokens.DIRECTED : Tokens.UNDIRECTED;
        	vals[1] = g.getNodeKeyField();
        	xml.start(Tokens.GRAPH, attr, vals, 2);
        }
        
        // print the nodes
        xml.comment("nodes");
        Iterator nodes = g.nodes();
        while ( nodes.hasNext() ) {
            Node tuple = (Node)nodes.next();
            if ( g.getNodeTable().getColumnCount() > 0 ) {            	
                xml.start(Tokens.NODE, Tokens.ID, String.valueOf(tuple.getRow()));
                for ( int i=0; i<g.getNodeTable().getColumnCount(); ++i ) {
                    String field = g.getNodeTable().getColumnName(i);                    
                    if (tuple.get(field) == null) {
                    	if (tuple.getDefault(field) == null) {
                    		continue;
                    	}
                    }
                    else if (tuple.get(field).equals(tuple.getDefault(field))) {
                    	continue;
                    }
                    DataParser dp = g.getNodeTable().getColumn(i).getParser();
                    String content = dp.format(tuple.get(i));
                    if (tuple.get(field) == null) {
                    	content = "";
                    }
                    xml.contentTag(Tokens.DATA, Tokens.KEY, field, content);
                }
                xml.end();
            } else {
                xml.tag(Tokens.NODE, Tokens.ID, String.valueOf(tuple.getRow()));
            }
        }
        // add a blank line
        xml.println();
        
        // print the edges
        String[] attr = new String[]{Tokens.ID, Tokens.SOURCE, Tokens.TARGET};
        String[] vals = new String[3];
        xml.comment("edges");
        Iterator edges = g.edges();
        while ( edges.hasNext() ) {
            Edge tuple = (Edge)edges.next();
            vals[0] = String.valueOf(tuple.getRow());
            vals[1] = String.valueOf(tuple.getSourceNode().getRow());
            vals[2] = String.valueOf(tuple.getTargetNode().getRow());
            
            if ( g.getEdgeTable().getColumnCount() > 2 ) {
                xml.start(Tokens.EDGE, attr, vals, 3);
                for ( int i=0; i<g.getEdgeTable().getColumnCount(); ++i ) {
                    String field = g.getEdgeTable().getColumnName(i);
                    DataParser dp = g.getEdgeTable().getColumn(i).getParser();
                    if (field.equals(g.getEdgeSourceField()) ||
                        field.equals(g.getEdgeTargetField())){
                    	if (g.getNodeKeyField() != null) {
                    		String str = dp.format(tuple.get(i));
                    		xml.contentTag(Tokens.DATA, Tokens.KEY, field, str);
						}
                   		continue;
                    }
//                    if ( field.equals(g.getEdgeSourceField()) ||
//                         field.equals(g.getEdgeTargetField()) )
//                        continue;
                    if (tuple.get(field) == null) {
                    	if (tuple.getDefault(field) == null){
                    		continue;
                    	}
					}
                    else if (tuple.get(field).equals(tuple.getDefault(field))) {
                    	continue;
                    }
                    String str = dp.format(tuple.get(i));
                    xml.contentTag(Tokens.DATA, Tokens.KEY, field, str);
                }
                xml.end();
            } else {
                xml.tag(Tokens.EDGE, attr, vals, 3);
            }
        }
        xml.end();

        // finish writing file
        xml.finish("</"+Tokens.GRAPHML+">\n");
        
        restoreParsers(g.getNodeTable(), nodeOldParser);
        restoreParsers(g.getEdgeTable(), edgeOldParser);
	}
	
    /**
     * Print a table schema to a GraphML file
     * @param xml the XMLWriter to write to
     * @param group the data group (node or edge) for the schema
     * @param t the schema
     */
	protected void printSchema(XMLWriter xml, String group, Table t, 
			                   String[] ignore) {
		String[] attr = new String[] {Tokens.ID, Tokens.FOR,
				Tokens.ATTRNAME, Tokens.ATTRTYPE };
		String[] vals = new String[4];

	OUTER:
		for ( int i=0; i<t.getColumnCount(); ++i ) {
			vals[0] = t.getColumnName(i);

			for ( int j=0; ignore!=null && j<ignore.length; ++j ) {
				if ( vals[0].equals(ignore[j]) )
					continue OUTER;
			}

			vals[1] = group;
			vals[2] = vals[0];
			vals[3] = formatType(t.getColumn(i));
			
			Object dflt = t.getDefault(t.getColumnName(i));
			if ( dflt == null ) {
				xml.tag(Tokens.KEY, attr, vals, 4);
			} else {
				xml.start(Tokens.KEY, attr, vals, 4);
				String dfltString = t.getColumn(i).getParser().format(dflt);
				xml.contentTag(Tokens.DEFAULT, dfltString);
				xml.end();
			}
		}
	}
    private String formatType(Column column){
    	DataParser parser = column.getParser();
    	Class type = parser.getType();
    	if (TYPES.containsKey(type)) {
			return TYPES.get(type);
		}
    	if (parser instanceof CollectionParser){
    		String result = "";
    		result += className(parser.getType());
    		while (parser instanceof CollectionParser){
    			parser = ((CollectionParser)parser).getElementParser();
    			if (parser == null) {
					return result;
				}
    			result += ":";
    			if (TYPES.containsKey(parser.getType()))
    				result += TYPES.get(parser.getType());
    			else
    				result += className(parser.getType());
    		}
    		return result;
    	}
    	return className(parser.getType());
    }
    private String className(Class type){
    	return type.toString().substring(type.toString().indexOf(" ")+1);
    }
}